package com.zhang.req.keyedProcess;

import com.zhang.second.utils.ClickEvent;
import com.zhang.second.utils.ClickSource;
import com.zhang.second.utils.UrlViewCountPerWindow;
import org.apache.flink.api.common.state.MapState;
import org.apache.flink.api.common.state.MapStateDescriptor;
import org.apache.flink.api.common.typeinfo.Types;
import org.apache.flink.configuration.Configuration;
import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
import org.apache.flink.streaming.api.functions.KeyedProcessFunction;
import org.apache.flink.util.Collector;

import java.util.ArrayList;
import java.util.List;

/**
 * 每隔5s输出过去10s的url访问情况
 * @title: 使用keyedProcess实现滑动窗口
 * @author: zhang
 * @date: 2022/2/3 10:13
 */
public class SlidingWindowInKeyedProcess {
    public static void main(String[] args) throws Exception {
        StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();
        env.setParallelism(1);

        env
                .addSource(new ClickSource())
                .keyBy(r -> r.url)
                .process(new KeyedProcessFunction<String, ClickEvent, UrlViewCountPerWindow>() {
                    private MapState<Long, List<ClickEvent>> mapState;
                    private long slidingWindowSize = 10 * 1000L;
                    private long slidingWindowStep = 5 * 1000L;

                    @Override
                    public void open(Configuration parameters) throws Exception {
                        mapState = getRuntimeContext().getMapState(
                                new MapStateDescriptor<Long, List<ClickEvent>>(
                                        "map-state",
                                        Types.LONG, Types.LIST(Types.POJO(ClickEvent.class))
                                )
                        );
                    }

                    @Override
                    public void processElement(ClickEvent value, KeyedProcessFunction<String, ClickEvent, UrlViewCountPerWindow>.Context ctx, Collector<UrlViewCountPerWindow> out) throws Exception {
                        long currTs = ctx.timerService().currentProcessingTime();
                        //计算属于的窗口开始时间
                        ArrayList<Long> arr = new ArrayList<>((int) (slidingWindowSize / slidingWindowStep));
                        long lastStart = getWindowStart(currTs, slidingWindowStep);
                        for (long start = lastStart; start > currTs - slidingWindowSize; start -= slidingWindowStep) {
                            arr.add(start);
                        }

                        for (Long startTime : arr) {
                            if (!mapState.contains(startTime)) {
                                ArrayList<ClickEvent> events = new ArrayList<>();
                                events.add(value);
                                mapState.put(startTime, events);
                            } else {
                                mapState.get(startTime).add(value);
                            }
                            ctx.timerService().registerProcessingTimeTimer(startTime + slidingWindowSize - 1L);
                        }
                    }

                    @Override
                    public void onTimer(long timestamp, KeyedProcessFunction<String, ClickEvent, UrlViewCountPerWindow>.OnTimerContext ctx, Collector<UrlViewCountPerWindow> out) throws Exception {
                        long windowEndTime = timestamp + 1L;
                        long windowStartTime = windowEndTime - slidingWindowSize;
                        String key = ctx.getCurrentKey();
                        long count = mapState.get(windowStartTime).size();
                        out.collect(new UrlViewCountPerWindow(
                                key,
                                count,
                                windowStartTime,
                                windowEndTime
                        ));
                        mapState.remove(windowStartTime);
                    }
                })
                .print();

        env.execute();
    }

    public static long getWindowStart(long timestamp, long step) {
        return timestamp - (timestamp + step) % step;
    }
}
